#################################################################
#		Assignment: OneStepWGCNA
#		Author:	Shawn Wang
#		Date: May 14, 2020
#       Updata: Jun 28, 2020
#       Version: 2
#################################################################

##=========updata info=================
# getopt was used as the args description. remove some unfriendly settings.
# suppressMessages to hide the package info
##========================Param setting==========================
library(getopt)
## args setting
command=matrix(c(
  'help', 'h', 0, 'logic', 'help information',
  'readcount', 'r', 1, 'character', 'inputfile: readcount matrix, geneID in row, sample names in column',
  'traitData', 't', 1, 'character', 'inputfile: trait data, if your data is quantitative traits, set sample names in row, trait in column, if it is discrete trait, set sample name in 1st column, levels in 2nd column',
  'RcCutoff', 'c', '2', 'integer', 'Noise remove: based on WGCNA FAQ, background noise should be removed, set the readcount value cutoff.(default = 6)',
  'samplePerc', 'p', '2', 'double', 'Noise remove: At least how many samples have readcount value greater than cutoff,(Range 0-1, default = 0.8)',
  'RemainGeneNum', 'g', '2', 'integer', 'WGCNA gene input: After removing the background noise, the gene needs to be screened a second time, and the gene with the highest MAD value among the samples will be retained.(default = 15000)',
  'output', 'o', 2, 'character', 'Output: project name for output file names.(default = system date)',
  'workdir', 'w', 1, 'character', 'working directory, you have type -w $PWD or -w `pwd`'
),byrow = T, ncol = 5)
args = getopt(command)
## help information
if (!is.null(args$help)) {
  cat(paste(getopt(command, usage = T), "\n"))
#  q(status=1)
}

## default value
if (is.null(args$RcCutoff)){
  args$RcCutoff = 6
}

if (is.null(args$samplePerc)){
  args$samplePerc = 0.8
}

if (is.null(args$RemainGeneNum)){
  args$RemainGeneNum = 15000
}

if (is.null(args$output)){
  args$output <- Sys.Date()
}

## Set working direction
wd <- args$workdir
setwd(wd)
##====================Section1.Prepration======================
## name args
readcount <- args$readcount
traitData <- args$traitData
RcCutoff <- args$RcCutoff
samplePerc <- args$samplePerc
RemainGeneNum <- args$RemainGeneNum
title <- args$output
## 02.load package
suppressMessages(library(DESeq2))
suppressMessages(library(ggplot2))
suppressMessages(library(dplyr))
suppressMessages(library(WGCNA))
suppressMessages(library(stringr))
suppressMessages(library(ape))
suppressMessages(library(reshape2))
options(stringsAsFactors = F)
#enableWGCNAThreads() ## for server
## 03.load 
## 3.1 readcount
rawcount <- read.delim(readcount, header = T,
                       sep = "\t")
rawcount <- data.frame(row.names = rawcount[,1],
                       rawcount[,-1])
## 3.2 trait
trait <- read.delim(traitData,header = T,
                    sep = "\t")
if (ncol(trait) == 2) {
  x <- trait
  Tcol = as.character(unique(x[,2]))
  b <- list()
  for (i in 1:length(Tcol)) {
    b[[i]] = data.frame(row.names = x[,1],
                        levels = ifelse(x[,2] == Tcol[i],1,0))
  }
  c <- bind_cols(b)
  c <- data.frame(row.names = x$name,
                  c)
  colnames(c) = Tcol
  pheTmp <- c
} else {
  pheTmp = data.frame(row.names = trait[,1],
                      trait[,-1])
}
## 04.data cleaning
## 4.1 condition
samnum <- ncol(rawcount)
casenum = ceiling(samnum/2)
controlnum = samnum - casenum
condition <- factor(c(rep("case",casenum),rep("control",controlnum)),levels = c("case","control"))
colData <- data.frame(row.names = colnames(rawcount), condition)
## remove background noise
x <- rawcount[apply(rawcount,1,function(x) sum(x > RcCutoff) > (samplePerc*ncol(rawcount))),]
dim(x)
## readcount standardization by DESeq2
dds <- DESeqDataSetFromMatrix(x, colData, design = ~ condition)
dds <- DESeq(dds)
vsd <- assay(varianceStabilizingTransformation(dds))
##====================Section2.WGCNA======================
## 01.parameter
datExpr = data.frame(vsd)
dim(datExpr)
type = "unsigned"
corType = "pearson"
corFnc = ifelse(corType=="pearson", cor, bicor)
maxPOutliers = ifelse(corType=="pearson",1,0.05)
robustY = ifelse(corType=="pearson",T,F)
Title = title
## 1.1 functions
source("/public/home/wangxiao/03.MyScript/MyScript/R/01.WGCNA/11.02.WGCNA.SFT.R")
source("/public/home/wangxiao/03.MyScript/MyScript/R/01.WGCNA/11.03.WGCNA.module.R")
source("/public/home/wangxiao/03.MyScript/MyScript/R/01.WGCNA/11.04.WGCNA.moduleTrait.R")
source("/public/home/wangxiao/03.MyScript/MyScript/R/01.WGCNA/11.05.WGCNA.HubGene.R")
## 02. sft
## 2.1 set remain gene number
GNC = 1-RemainGeneNum/nrow(datExpr) 
## 2.2 sft calculate
WGCNA.SFT(Title = title,
          datExpr = datExpr,
          GeneNumCut = GNC)

print("#################==01.SFT calculation finish!==###########################")

## confirm power
if (is.na(power)){
  power = ifelse(nSamples<20, ifelse(type == "unsigned", 9, 18),
                  ifelse(nSamples<30, ifelse(type == "unsigned", 8, 16),
                         ifelse(nSamples<40, ifelse(type == "unsigned", 7, 14),
                               ifelse(type == "unsigned", 6, 12))
                 )
  )
}
print(paste("The most suitable Power:",power))

## 2.2 OneStepNetwork
WGCNA.oneStepNetWork(Title = Title)

print("#################==02.Network has bean constructed!==###########################")

# 2.3 Module-Trait relation
phenotype = pheTmp
WGCNA.ModuleTrait(Title = Title,
phenotype = phenotype)
save(MEs_col,nSamples,corType,file = "datatraitbase.Rdata")
print("#################==03.ModuleTraits value has finish!==###########################")

## 2.4 IntramodularConnectivity
Title = title
connet=abs(cor(datExpr,use="p"))^6
Alldegrees1=intramodularConnectivity(connet, moduleColors)
###(3) Generalizing intramodular connectivity for all genes on the array
datKME=signedKME(datExpr, MEs_col, outputColumnName="MM.")
write.table(datKME, paste(Title,"Conectivity_of_each_modular.xls",sep = "."),
            sep = "\t",
            row.names = T,
            quote = F)
save(... = datExpr, power,moduleColors,file = paste(Title,".Modular.Rdata",sep = ""))
print("#################==All work finish!==###########################")
